"""
erosion_jet.py

此模块演示了如何针对 Jet 图中高值和低值两种“前景色”进行形态学腐蚀操作。
Jet 图是一种常用的颜色映射图，本模块通过阈值分割将其分为高值前景、低值前景和中间背景，
然后分别对高值和低值前景的二值掩码执行腐蚀操作，最后将结果合并并可视化。

主要功能包括：
1. 创建模拟的 Jet 图
2. 定义阈值来区分不同区域
3. 创建高值和低值前景的二值掩码
4. 定义结构元素并执行腐蚀操作
5. 合并腐蚀后的结果
6. 可视化原始图像、掩码、腐蚀结果和合并后的图像
"""


import cv2
import numpy as np
import matplotlib.pyplot as plt

def _create_jet_image():
    """
    创建一个模拟的 Jet 图，包含高值区域、低值区域和中间背景区域。
    """
    jet_image = np.full((15, 15), 127, dtype=np.uint8)  # 初始设置为中间色调 (背景)
    # 添加高值“前景”区域 (接近255)
    jet_image[2:6, 2:6] = 250  # 大的高值方块
    jet_image[0, 14] = 255  # 小的高值噪声点
    # 添加低值“前景”区域 (接近0)
    jet_image[9:13, 9:13] = 5  # 大的低值方块
    jet_image[14, 0] = 0  # 小的低值噪声点
    print("原始 Jet 图：")
    print(jet_image)
    return jet_image

def _create_thresholds():
    """
    定义阈值来区分高值、低值前景色和中间背景。
    """
    lower_threshold_for_low_foreground = 50
    upper_threshold_for_high_foreground = 200
    return lower_threshold_for_low_foreground, upper_threshold_for_high_foreground

def _create_binary_masks(jet_image, lower_threshold, upper_threshold):
    """
    创建高值和低值前景的二值掩码。
    """
    _, high_foreground_mask = cv2.threshold(jet_image, upper_threshold, 255, cv2.THRESH_BINARY)
    print("\n高值前景掩码 (白色表示高值前景):")
    print(high_foreground_mask)
    _, low_foreground_mask = cv2.threshold(jet_image, lower_threshold, 255, cv2.THRESH_BINARY_INV)
    print("\n低值前景掩码 (白色表示低值前景):")
    print(low_foreground_mask)
    return high_foreground_mask, low_foreground_mask

def _create_kernel():
    """
    定义一个 3x3 的矩形结构元素。
    """
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
    print("\n使用的结构元素 (3x3 矩形):")
    print(kernel)
    return kernel

def _perform_erosion(high_foreground_mask, low_foreground_mask, kernel):
    """
    对高值和低值前景的二值掩码分别执行腐蚀操作。
    """
    eroded_high_foreground_mask = cv2.erode(high_foreground_mask, kernel, iterations=1)
    eroded_low_foreground_mask = cv2.erode(low_foreground_mask, kernel, iterations=1)
    print("\n腐蚀后的高值前景掩码：")
    print(eroded_high_foreground_mask)
    print("\n腐蚀后的低值前景掩码：")
    print(eroded_low_foreground_mask)
    return eroded_high_foreground_mask, eroded_low_foreground_mask

def _combine_eroded_images(jet_image, eroded_high_foreground_mask, eroded_low_foreground_mask):
    """
    将腐蚀后的结果合并回一张图像。
    """
    combined_eroded_image = np.full(jet_image.shape, 127, dtype=np.uint8)  # 默认为中间背景色
    combined_eroded_image[eroded_high_foreground_mask == 255] = 250  # 将腐蚀后的高值前景设为高值
    combined_eroded_image[eroded_low_foreground_mask == 255] = 5  # 将腐蚀后的低值前景设为低值
    print("\n合并后的腐蚀图像：")
    print(combined_eroded_image)
    return combined_eroded_image

def _visualize_results(jet_image, high_foreground_mask, low_foreground_mask, eroded_high_foreground_mask, eroded_low_foreground_mask, combined_eroded_image, lower_threshold, upper_threshold):
    """
    可视化原始图像、掩码、腐蚀后的掩码和合并后的图像。
    """
    fig, axes = plt.subplots(2, 3, figsize=(12, 8))
    axes[0, 0].imshow(jet_image, cmap='jet', vmin=0, vmax=255)  # 使用 'jet' colormap
    axes[0, 0].set_title('Original Jet Image')
    axes[0, 0].axis('off')
    axes[0, 1].imshow(high_foreground_mask, cmap='gray', vmin=0, vmax=255)
    axes[0, 1].set_title(f'High Foreground Mask (> {upper_threshold})')
    axes[0, 1].axis('off')
    axes[0, 2].imshow(low_foreground_mask, cmap='gray', vmin=0, vmax=255)
    axes[0, 2].set_title(f'Low Foreground Mask (< {lower_threshold})')
    axes[0, 2].axis('off')
    axes[1, 0].imshow(eroded_high_foreground_mask, cmap='gray', vmin=0, vmax=255)
    axes[1, 0].set_title('Eroded High Foreground Mask')
    axes[1, 0].axis('off')
    axes[1, 1].imshow(eroded_low_foreground_mask, cmap='gray', vmin=0, vmax=255)
    axes[1, 1].set_title('Eroded Low Foreground Mask')
    axes[1, 1].axis('off')
    axes[1, 2].imshow(combined_eroded_image, cmap='jet', vmin=0, vmax=255)
    axes[1, 2].set_title('Combined Eroded Image (Jet)')
    axes[1, 2].axis('off')
    plt.tight_layout()
    plt.show()


if __name__ == "__main__":
    """
    演示如何针对 Jet 图中高值和低值两种“前景色”进行腐蚀操作。
    """
    # STEP1：创建 Jet 图
    jet_image = _create_jet_image()

    # STEP2：创建阈值
    lower_threshold, upper_threshold = _create_thresholds()
    
    # STEP3：创建二值掩码
    high_foreground_mask, low_foreground_mask = _create_binary_masks(jet_image, lower_threshold, upper_threshold)
    
    # STEP4：创建结构元素并用此结构执行腐蚀操作
    kernel = _create_kernel()
    eroded_high_foreground_mask, eroded_low_foreground_mask = _perform_erosion(high_foreground_mask, low_foreground_mask, kernel)

    # STEP5：合并腐蚀后的结果
    combined_eroded_image = _combine_eroded_images(jet_image, eroded_high_foreground_mask, eroded_low_foreground_mask)

    # STEP6：可视化结果
    _visualize_results(jet_image, high_foreground_mask, low_foreground_mask, eroded_high_foreground_mask, eroded_low_foreground_mask, combined_eroded_image, lower_threshold, upper_threshold)
